package com.example.edu.Utils;

import cn.hutool.json.JSONUtil;
import com.example.edu.Dto.PromptDto;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.rxjava3.core.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static com.example.edu.Utils.MyConstants.AI_API_KEY;

@Component
public class AIUtil {
    private static final Logger logger = LoggerFactory.getLogger(AIUtil.class);

    public String getAnswer(PromptDto promptData,int is_json,String model_code){
        ClientV4 client = new ClientV4.Builder(AI_API_KEY)
                .networkConfig(480,30,240,180, TimeUnit.SECONDS)
                .build();
        List<ChatMessage> messages = generateMessages(promptData);
        List<ChatTool> chatTools = generateWebSearch(promptData);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(model_code)
                .invokeMethod(Constants.invokeMethod)
                .stream(false)
                .messages(messages)
                .tools(chatTools)
                .build();
        try {
            ModelApiResponse modelApiResponse = client.invokeModelApi(chatCompletionRequest);
            Object result = modelApiResponse.getData().getChoices().get(0).getMessage().getContent();
            String r = null;
            if (is_json == 1) {
                r = JSONUtil.toJsonStr(result);
                // 移除换行符、回车符等不可见字符
                r = r.replaceFirst("```json\\n", "").replace("```", "");
            }
            else
                r = result.toString();
            return r;
        } catch (Exception e) {
            logger.error(e.getMessage());
            return null;
        }
    }

    public Flowable<String> getAnswerSSE(PromptDto promptData){
        ClientV4 client = new ClientV4.Builder(AI_API_KEY)
                .networkConfig(480,30,240,180, TimeUnit.SECONDS)
                .build();
        List<ChatMessage> messages = generateMessages(promptData);
        List<ChatTool> chatTools = generateWebSearch(promptData);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4Plus)
                .invokeMethod(Constants.invokeMethod)
                .stream(true)
                .messages(messages)
                .tools(chatTools)
                .build();
        return getStringFlowable(client, chatCompletionRequest);
    }

    public Flowable<String> getAnswerSSE(PromptDto promptData,String chatId,String model_code){
        ClientV4 client = new ClientV4.Builder(AI_API_KEY)
                .networkConfig(480,30,240,180, TimeUnit.SECONDS)
                .build();
        List<ChatMessage> messages = generateMessages(promptData);
        List<ChatTool> chatTools = generateWebSearch(promptData);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(model_code)
                .invokeMethod(Constants.invokeMethod)
                .stream(true)
                .requestId(chatId)
                .messages(messages)
                .tools(chatTools)
                .build();
        return getStringFlowable(client, chatCompletionRequest);
    }

    private Flowable<String> getStringFlowable(ClientV4 client, ChatCompletionRequest chatCompletionRequest) {
        try{
            ModelApiResponse modelApiResponse = client.invokeModelApi(chatCompletionRequest);
            if(!modelApiResponse.isSuccess())
                return Flowable.error(new RuntimeException("Model API call failed"));
            return Flowable.fromPublisher(modelApiResponse.getFlowable().map(data -> {
                String content = data.getChoices().get(0).getDelta().getContent();
                return content != null ? content : "";
            }));
        } catch (Exception e) {
            logger.error(e.getMessage());
            return Flowable.error(new RuntimeException("Model API call failed"));
        }
    }

    private List<ChatMessage> generateMessages(PromptDto promptData){
        List<ChatMessage> messages = new ArrayList<>();
        if(promptData.getSystem_prompt() != null && !promptData.getSystem_prompt().isEmpty()){
            StringBuilder systemPrompt = new StringBuilder();
            for(String system_prompt : promptData.getSystem_prompt())
                systemPrompt.append(system_prompt).append("\n");
            messages.add(new ChatMessage(ChatMessageRole.SYSTEM.value(), systemPrompt.toString()));
        }
        if(promptData.getUser_prompt() != null && !promptData.getUser_prompt().isEmpty() && promptData.getAssistant_prompt() != null && !promptData.getAssistant_prompt().isEmpty()){
            for(int i = 0; i < promptData.getAssistant_prompt().size(); i++) {
                messages.add(new ChatMessage(ChatMessageRole.USER.value(), promptData.getUser_prompt().get(i)));
                messages.add(new ChatMessage(ChatMessageRole.ASSISTANT.value(),promptData.getAssistant_prompt().get(i)));
            }
            messages.add(new ChatMessage(ChatMessageRole.USER.value(), promptData.getUser_prompt().get(promptData.getUser_prompt().size()-1)));
        }
        else if(promptData.getUser_prompt() != null && !promptData.getUser_prompt().isEmpty()){
            StringBuilder userPrompt = new StringBuilder();
            for(String user_prompt : promptData.getUser_prompt())
                userPrompt.append(user_prompt);
            messages.add(new ChatMessage(ChatMessageRole.USER.value(), userPrompt.toString()));
        }
        if(promptData.getFunction_prompt() != null && !promptData.getFunction_prompt().isEmpty()){
            for(String function_prompt : promptData.getFunction_prompt())
                messages.add(new ChatMessage(ChatMessageRole.FUNCTION.value(), function_prompt));
        }
        return messages;
    }

    private List<ChatTool> generateWebSearch(PromptDto promptData){
        List<ChatTool> chatTools = new ArrayList<>();
        if (promptData.getWeb_search_prompt() != null){
            ChatTool chatTool = new ChatTool();
            chatTool.setType(ChatToolType.WEB_SEARCH.value());
            WebSearch webSearch = new WebSearch();
            webSearch.setSearch_query(promptData.getWeb_search_prompt());
            webSearch.setEnable(Boolean.TRUE);
            chatTool.setWeb_search(webSearch);
            chatTools.add(chatTool);
        }
        return chatTools;
    }
}
